/*
 * Environment:
 * Microsoft Visual Studio 2008 Version 9.0.21022.8 RTM
 * Microsoft (R) 32-bit C/C++ Optimizing Compiler Version 15.00.21022.08 for 80x86
 *
 * Authors: Sergey Zamyslov, Sergey Moshkov
 */

#include <stdio.h>

#define DEBUG false

#define THREADS_PER_BLOCK 256
#define ROWS_PER_MATRIX_BLOCK 32 // >= 2.

float function(float x) {
    return cosf(x) / x;
}

void print(float* data, int n)
{
    int i;

    n = min(n, 12);

    for (i = 0; i < n; i++)
    {
        printf("%11.4f", data[i]);
    }
}

void initializeKnotsAndNodes(float a, float b,
                             float* knots, int knotCount,
                             float* nodesX, float* nodesY)
{
    int    i;
    float step = (b - a) / (knotCount - 1);

    float d = step * 0.9f;
    float dOver2 = d / 2.0f;

    knots[0]             = a;
    knots[knotCount - 1] = b;

    nodesX[0] = knots[0];
    nodesY[0] = function(nodesX[0]);

    srand((int)time(NULL));

    for (i = 1; i < knotCount - 1; i++)
    {
        knots[i] = i * step;
        knots[i] += (((float)rand() / RAND_MAX) * d) - dOver2;

        nodesX[i] = (knots[i] + knots[i - 1]) / 2;
        nodesY[i] = function(nodesX[i]);
    }

    i = knotCount - 1;
    nodesX[i] = (knots[i] + knots[i - 1]) / 2;
    nodesY[i] = function(nodesX[i]);

    nodesX[knotCount] = knots[knotCount - 1];
    nodesY[knotCount] = function(nodesX[knotCount]);
}

void solveTDM(int equationCount,
              float* subDiagonal, float* mainDiagonal, float* supDiagonal,
              float* b,
              float* x)
{
    // Initialization
    int    i;
    float* mainDiagonalWork = new float[equationCount],
         * bWork            = new float[equationCount];

	mainDiagonalWork[0] = mainDiagonal[0];
	bWork[0]            = b[0];

    // Forward
    for (i = 1; i < equationCount; i++)
    {
        float m = subDiagonal[i-1] / mainDiagonalWork[i-1];

        mainDiagonalWork[i] = mainDiagonal[i] - m * supDiagonal[i-1];
        bWork[i]            = b[i] - m * bWork[i-1];
	}

    // Backward
    x[equationCount - 1] = bWork[equationCount - 1] / mainDiagonalWork[equationCount - 1];	
    for (i = equationCount - 2; i >= 0; i--)
    {
        x[i] = (bWork[i] - supDiagonal[i] * x[i + 1]) / mainDiagonalWork[i];
    }

    delete[] mainDiagonalWork;
    delete[] bWork;
}

__global__ void solveTDM_BlockForward(int equationCount,
                                      float* subDiagonal, float* mainDiagonal, float* supDiagonal,
                                      float* b)
{
    // Matrix block (one thread handles one matrix block).
    int blockIndex = blockDim.x * blockIdx.x + threadIdx.x;

    int blockStartRow = blockIndex * ROWS_PER_MATRIX_BLOCK;
    int blockEndRow = blockStartRow + ROWS_PER_MATRIX_BLOCK - 1;

    float factor;

    for(int blockRowIndex = blockStartRow; blockRowIndex < blockEndRow; ++blockRowIndex)
    {
        factor = subDiagonal[blockRowIndex] / mainDiagonal[blockRowIndex];

        if(blockIndex > 0)
        {
            // Non-diagonal elements.
            subDiagonal[blockRowIndex] = -subDiagonal[blockRowIndex - 1] * factor;
        }
        else
        {
            subDiagonal[blockRowIndex] = 0.0;
        }

        mainDiagonal[blockRowIndex + 1] -= supDiagonal[blockRowIndex] * factor;
        b[blockRowIndex + 1] -= b[blockRowIndex] * factor;
    }
}

__global__ void solveTDM_BlockBackward(int equationCount,
                                       float* subDiagonal, float* mainDiagonal, float* supDiagonal,
                                       float* b)
{
    int blockIndex = blockDim.x * blockIdx.x + threadIdx.x;

    int blockStartRow = blockIndex * ROWS_PER_MATRIX_BLOCK;
    int blockEndRow = blockStartRow + ROWS_PER_MATRIX_BLOCK - 1;
    blockStartRow = max(blockStartRow, 1);

    float factor;

    for(int blockRowIndex = blockEndRow - 1; blockRowIndex >= blockStartRow; --blockRowIndex)
    {
        factor = supDiagonal[blockRowIndex - 1] / mainDiagonal[blockRowIndex];

        if(blockIndex > 0)
        {
            if(blockRowIndex > blockStartRow)
            {
                subDiagonal[blockRowIndex - 2] -= subDiagonal[blockRowIndex - 1] * factor;
            }
            else
            {
                // If we edit element out of block, we take element from main diagonal.
                mainDiagonal[blockRowIndex - 1] -= subDiagonal[blockRowIndex - 1] * factor;
            }
        }

        // Non-diagonal elements.
        supDiagonal[blockRowIndex - 1] = -supDiagonal[blockRowIndex] * factor;

        b[blockRowIndex - 1] -= b[blockRowIndex] * factor;
    }
}

__global__ void getAuxiliaryMatrixCoeffs(int matrixBlockCount,
                                         float* subDiagonal, float* mainDiagonal, float* supDiagonal,
                                         float* b,
                                         float* auxiliaryMatrix)
{
    int blockIndex = blockDim.x * blockIdx.x + threadIdx.x;

    int blockStartRow = blockIndex * ROWS_PER_MATRIX_BLOCK;
    int blockEndRow = blockStartRow + ROWS_PER_MATRIX_BLOCK - 1;

    if(blockIndex > 0)
    {
        auxiliaryMatrix[0 * matrixBlockCount + blockIndex - 1] = subDiagonal[blockEndRow - 1];
    }

    if(blockIndex < matrixBlockCount - 1)
    {
        auxiliaryMatrix[2 * matrixBlockCount + blockIndex] = supDiagonal[blockEndRow];
    }

    auxiliaryMatrix[1 * matrixBlockCount + blockIndex] = mainDiagonal[blockEndRow];
    auxiliaryMatrix[3 * matrixBlockCount + blockIndex] = b[blockEndRow];
}

__global__ void setX(float* auxiliaryX,
                     float* x)
{
    int blockIndex = blockDim.x * blockIdx.x + threadIdx.x;

    int blockStartRow = blockIndex * ROWS_PER_MATRIX_BLOCK;
    int blockEndRow = blockStartRow + ROWS_PER_MATRIX_BLOCK - 1;

    if(DEBUG)
    {
        for(int blockRowIndex = blockStartRow; blockRowIndex < blockEndRow; ++blockRowIndex)
        {
            x[blockRowIndex] = 0.0f;
        }
    }

    x[blockEndRow] = auxiliaryX[blockIndex];
}

__global__ void calculateX(int equationCount,
                           float* subDiagonal, float* mainDiagonal, float* supDiagonal,
                           float* b,
                           float* x)
{
    int blockIndex = blockDim.x * blockIdx.x + threadIdx.x;

    int blockStartRow = blockIndex * ROWS_PER_MATRIX_BLOCK;
    int blockEndRow = blockStartRow + ROWS_PER_MATRIX_BLOCK - 1;

    // Use local variables and not arrays.
    float xEndRow = x[blockEndRow];
    float xStartRow;
    if(blockIndex > 0)
    {
        xStartRow = x[blockStartRow - 1];
    }

    float xRow;
    for(int blockRowIndex = blockStartRow; blockRowIndex < blockEndRow; ++blockRowIndex)
    {
        xRow = b[blockRowIndex];
        if(blockIndex > 0)
        {
            xRow -= subDiagonal[blockRowIndex - 1] * xStartRow;
        }
        xRow -= supDiagonal[blockRowIndex] * xEndRow;
        xRow /= mainDiagonal[blockRowIndex];

        x[blockRowIndex] = xRow;
    }
}

double solveTDM_Block(int equationCount,
                      float* subDiagonal, float* mainDiagonal, float* supDiagonal,
                      float* b,
                      float* x)
{
    // Timing.
	clock_t start, finish;
	double time;

    cudaError_t error;

    // Check input.
    if(equationCount % ROWS_PER_MATRIX_BLOCK != 0)
    {
        printf("equationCount cannot be divided evenly by %d (ROWS_PER_MATRIX_BLOCK).", ROWS_PER_MATRIX_BLOCK);
        exit(-1);
    }

    int matrixBlockCount = equationCount / ROWS_PER_MATRIX_BLOCK;

    if(matrixBlockCount % THREADS_PER_BLOCK != 0)
    {
        printf("matrixBlockCount cannot be divided evenly by %d (THREADS_PER_BLOCK).", THREADS_PER_BLOCK);
        exit(-1);
    }

    if(DEBUG)
    {
        printf("subDiagonal:\n");
        print(subDiagonal, equationCount - 1);
        printf("\n");

        printf("mainDiagonal:\n");
        print(mainDiagonal, equationCount);
        printf("\n");

        printf("supDiagonal:\n");
        print(supDiagonal, equationCount - 1);
        printf("\n");

        printf("b:\n");
        print(b, equationCount);
        printf("\n\n");
    }

    start = clock();

    // Allocate device memory.
    size_t dataElementSize = sizeof(float);

    float* deviceSubDiagonal = 0;
	float* deviceMainDiagonal = 0;
    float* deviceSupDiagonal = 0;
    float* deviceB = 0;
    float* deviceX = 0;
    float* deviceAuxiliaryMatrix = 0;

    cudaMalloc(&deviceSubDiagonal, dataElementSize * (equationCount - 1));
    cudaMalloc(&deviceMainDiagonal, dataElementSize * equationCount);
    cudaMalloc(&deviceSupDiagonal, dataElementSize * (equationCount - 1));
    cudaMalloc(&deviceB, dataElementSize * equationCount);
    cudaMalloc(&deviceX, dataElementSize * equationCount);
    cudaMalloc(&deviceAuxiliaryMatrix, dataElementSize * matrixBlockCount * 4);

    // Check for errors.
	error = cudaGetLastError();
    if(error == cudaSuccess)
    {
        // Copy data to device.
        cudaMemcpy(deviceSubDiagonal, subDiagonal,
            dataElementSize * (equationCount - 1), cudaMemcpyHostToDevice);

        cudaMemcpy(deviceMainDiagonal, mainDiagonal,
            dataElementSize * equationCount, cudaMemcpyHostToDevice);

        cudaMemcpy(deviceSupDiagonal, supDiagonal,
            dataElementSize * (equationCount - 1), cudaMemcpyHostToDevice);

        cudaMemcpy(deviceB, b,
            dataElementSize * equationCount, cudaMemcpyHostToDevice);



        // CUDA kernels.
        int cudaBlockCount = matrixBlockCount / THREADS_PER_BLOCK;

        if(DEBUG)
        {
            printf("Equation count = %d.\n", equationCount);
            printf("CUDA block count = %d.\n", cudaBlockCount);
            printf("Threads per block = %d.\n", THREADS_PER_BLOCK);
            printf("Rows per thread = %d.\n", ROWS_PER_MATRIX_BLOCK);
            printf("\n");
        }

        solveTDM_BlockForward<<<cudaBlockCount, THREADS_PER_BLOCK>>>(
            equationCount,
            deviceSubDiagonal, deviceMainDiagonal, deviceSupDiagonal, deviceB);

        cudaThreadSynchronize();

        solveTDM_BlockBackward<<<cudaBlockCount, THREADS_PER_BLOCK>>>(
            equationCount,
            deviceSubDiagonal, deviceMainDiagonal, deviceSupDiagonal, deviceB);

        cudaThreadSynchronize();

        // Solve auxiliary matrix.
        getAuxiliaryMatrixCoeffs<<<cudaBlockCount, THREADS_PER_BLOCK>>>(
            matrixBlockCount,
            deviceSubDiagonal, deviceMainDiagonal, deviceSupDiagonal, deviceB,
            deviceAuxiliaryMatrix);

        float* auxiliaryMatrix[5];
        auxiliaryMatrix[0] = new float[matrixBlockCount - 1];
        auxiliaryMatrix[1] = new float[matrixBlockCount];
        auxiliaryMatrix[2] = new float[matrixBlockCount - 1];
        auxiliaryMatrix[3] = new float[matrixBlockCount];
        auxiliaryMatrix[4] = new float[matrixBlockCount];

        cudaThreadSynchronize();

        // Copy auxiliary matrix to host.
        cudaMemcpy(auxiliaryMatrix[0], &deviceAuxiliaryMatrix[0 * matrixBlockCount],
            dataElementSize * (matrixBlockCount - 1), cudaMemcpyDeviceToHost);

        cudaMemcpy(auxiliaryMatrix[1], &deviceAuxiliaryMatrix[1 * matrixBlockCount],
            dataElementSize * matrixBlockCount, cudaMemcpyDeviceToHost);

        cudaMemcpy(auxiliaryMatrix[2], &deviceAuxiliaryMatrix[2 * matrixBlockCount],
            dataElementSize * (matrixBlockCount - 1), cudaMemcpyDeviceToHost);

        cudaMemcpy(auxiliaryMatrix[3], &deviceAuxiliaryMatrix[3 * matrixBlockCount],
            dataElementSize * matrixBlockCount, cudaMemcpyDeviceToHost);

        if(DEBUG)
        {
            printf("Auxiliary subDiagonal:\n");
            print(auxiliaryMatrix[0], matrixBlockCount - 1);
            printf("\n");

            printf("Auxiliary mainDiagonal:\n");
            print(auxiliaryMatrix[1], matrixBlockCount);
            printf("\n");

            printf("Auxiliary supDiagonal:\n");
            print(auxiliaryMatrix[2], matrixBlockCount - 1);
            printf("\n");

            printf("Auxiliary b:\n");
            print(auxiliaryMatrix[3], matrixBlockCount);
            printf("\n\n");
        }

        solveTDM(matrixBlockCount,
            auxiliaryMatrix[0], auxiliaryMatrix[1], auxiliaryMatrix[2],
            auxiliaryMatrix[3],
            auxiliaryMatrix[4]);

        if(DEBUG)
        {
            printf("Auxiliary x:\n");
            print(auxiliaryMatrix[4], matrixBlockCount);
            printf("\n\n");
        }

        // Copy solution of auxiliary matrix back to device.
        cudaMemcpy(deviceAuxiliaryMatrix, auxiliaryMatrix[4],
            dataElementSize * matrixBlockCount, cudaMemcpyHostToDevice);

        setX<<<cudaBlockCount, THREADS_PER_BLOCK>>>(
            deviceAuxiliaryMatrix,
            deviceX);

        if(DEBUG)
        {
            cudaThreadSynchronize();

            cudaMemcpy(x, deviceX,
            dataElementSize * equationCount, cudaMemcpyDeviceToHost);

            printf("x on midway:\n");
            print(x, equationCount);
            printf("\n\n");
        }

        delete[] auxiliaryMatrix[0];
        delete[] auxiliaryMatrix[1];
        delete[] auxiliaryMatrix[2];
        delete[] auxiliaryMatrix[3];
        delete[] auxiliaryMatrix[4];

        cudaThreadSynchronize();

        calculateX<<<cudaBlockCount, THREADS_PER_BLOCK>>>(
            equationCount,
            deviceSubDiagonal, deviceMainDiagonal, deviceSupDiagonal, deviceB,
            deviceX);

        cudaThreadSynchronize();

        // Copy solution of matrix to host.
        cudaMemcpy(x, deviceX,
            dataElementSize * equationCount, cudaMemcpyDeviceToHost);

        // Check for errors.
        error = cudaGetLastError();

        if(DEBUG)
        {
            printf("x:\n");
            print(x, equationCount);
            printf("\n\n");
        }
    }

    // Free device memory.
	if(deviceSubDiagonal != 0) {
		cudaFree(deviceSubDiagonal);
		deviceSubDiagonal = 0;
	}

    if(deviceMainDiagonal != 0) {
		cudaFree(deviceMainDiagonal);
		deviceMainDiagonal = 0;
	}

    if(deviceSupDiagonal != 0) {
		cudaFree(deviceSupDiagonal);
		deviceSupDiagonal = 0;
	}

    if(deviceB != 0) {
		cudaFree(deviceB);
		deviceB = 0;
	}

    if(deviceX != 0) {
		cudaFree(deviceX);
		deviceX = 0;
	}

    if(deviceAuxiliaryMatrix != 0) {
		cudaFree(deviceAuxiliaryMatrix);
		deviceAuxiliaryMatrix = 0;
	}

    finish = clock();
    time = ((double)finish - (double)start) / CLOCKS_PER_SEC;

    if(error != cudaSuccess)
    {
        printf("Error occurred (%s).\n", cudaGetErrorString(error));
    }

    return time;
}

void calculateSplineCoeffs(float* knots, int knotCount, float* nodesY,
                           float* coeffs)
{
    // Prepare matrix.
    float* subDiagonal  = new float[knotCount - 1],
         * supDiagonal  = new float[knotCount - 1],
         * mainDiagonal = new float[knotCount],
         * b            = new float[knotCount];

    for (int i = 1; i < knotCount - 1; i++)
    {
        subDiagonal[i - 1] = knots[i] - knots[i - 1];
        supDiagonal[i - 1] = knots[i] - knots[i - 1];
        mainDiagonal[i]    = 3 * (knots[i + 1] - knots[i - 1]);
        b[i - 1]           = 8 * (nodesY[i] - nodesY[i - 1]);
    }

    mainDiagonal[0] = 3 * (knots[1] - knots[0]);
    mainDiagonal[knotCount - 1] = 3 * (knots[knotCount - 1] - knots[knotCount - 2]);
    subDiagonal[knotCount - 2]  = knots[knotCount - 1] - knots[knotCount - 2];
    supDiagonal[knotCount - 2]  = knots[knotCount - 1] - knots[knotCount - 2];
    b[knotCount - 2]            = 8 * (nodesY[knotCount - 1] - nodesY[knotCount - 2]);
    b[knotCount - 1]            = 8 * (nodesY[knotCount] - nodesY[knotCount - 1]);

    // Solve matrix.
    double time = solveTDM_Block(knotCount, subDiagonal, mainDiagonal, supDiagonal, b, coeffs);
    printf("Time: %f\n\n\n", time);

    delete[] subDiagonal;
    delete[] supDiagonal;
    delete[] mainDiagonal;
    delete[] b;
}

float getSplinePointY(int knotCount, float* knots,
                      float* nodesX, float* nodesY,
                      float* coeffs,
                      float x)
{
    int i;
    float c;

    for (i = 1; i < knotCount; ++i)
    {
        if (x <= knots[i])
        {
            break;
        }
    }

    if (i < knotCount)
    {
        c = x - nodesX[i];

        return nodesY[i] +
            0.5f * (coeffs[i] + coeffs[i - 1]) * c + 
            (coeffs[i] - coeffs[i - 1]) * c * c / (2 * (knots[i] - knots[i - 1]));
    }

    return 0.0f;
}

int main(int argc, char** argv) {

    // Interpolation parameters.
    float intervalStart = 0.5f;
    float intervalEnd = 10.0f;
    // Each thread operates with one matrix block.
    int size = 1024000 / (ROWS_PER_MATRIX_BLOCK * THREADS_PER_BLOCK);
    int knotCount = size * THREADS_PER_BLOCK * ROWS_PER_MATRIX_BLOCK;

    // Output parameters.
    bool output = true;
    bool outputNodes = false;
    char* outputFileName = "chart.csv";
    float chartDx = (intervalEnd - intervalStart) / 1000.0f;

    // Interpolation input data.
    float* knots  = new float[knotCount];
    float* nodesX = new float[knotCount + 1];
    float* nodesY = new float[knotCount + 1];

    // Interpolation output data.
    float* coeffs  = new float[knotCount];

    initializeKnotsAndNodes(intervalStart, intervalEnd, knots, knotCount, nodesX, nodesY);
    calculateSplineCoeffs(knots, knotCount, nodesY, coeffs);

    // Output.
    if(output)
    {
        FILE* outputFile = fopen(outputFileName, "w+");

        fprintf(outputFile, "X;Y;splineY;nodesY\n");
        float chartX;
        for(chartX = intervalStart; chartX < intervalEnd; chartX += chartDx) {
            fprintf(outputFile, "%f;%f;%f;NA\n",
                chartX,
                function(chartX),
                getSplinePointY(knotCount, knots, nodesX, nodesY, coeffs, chartX));
        }

        chartX = intervalEnd;
        fprintf(outputFile, "%f;%f;%f;NA\n", chartX,
            function(chartX),
            getSplinePointY(knotCount, knots, nodesX, nodesY, coeffs, chartX));

        if(outputNodes)
        {
            for(int i = 0; i <= knotCount; ++i) {
                fprintf(outputFile, "%f;NA;NA;%f\n",
                    nodesX[i],
                    nodesY[i]);
            }
        }

        fclose(outputFile);
    }

    delete[] coeffs;
    delete[] knots;
    delete[] nodesX;
    delete[] nodesY;

    return 0;
}
